package com.zjc.agent.ai.controller;

import com.zjc.agent.ai.service.AIFlightHelperService;
import com.zjc.agent.ai.utils.TokenContextHolder;
import jakarta.annotation.Resource;
import jakarta.servlet.http.HttpServletRequest;
import lombok.extern.slf4j.Slf4j;
import org.springframework.http.HttpRequest;
import org.springframework.http.codec.ServerSentEvent;
import org.springframework.web.bind.annotation.CrossOrigin;
import org.springframework.web.bind.annotation.RequestMapping;
import org.springframework.web.bind.annotation.RestController;
import org.springframework.web.servlet.mvc.method.annotation.SseEmitter;
import reactor.core.publisher.Flux;

import java.io.IOException;
import java.net.InetAddress;
import java.net.UnknownHostException;

/**
 * @Author: JinChen Zhu
 * @CreateTime: 2025-08-01  03:11
 */
@Slf4j
@RestController
@RequestMapping("/ai")
@CrossOrigin
public class AiController {
    @Resource
    private AIFlightHelperService aiFlightHelperService;

    /**
     * 流式聊天
     *
     * @param memoryId 内存ID
     * @param message  用户发送的消息
     * @return 聊天结果
     */
    @RequestMapping("/chat/stream")
    public Flux<ServerSentEvent<String>> chatStream(int memoryId, String message) {
        return aiFlightHelperService.chatStream(memoryId, message)
                .map(chunk -> ServerSentEvent.<String>builder()
                        .data(chunk)
                        .build())
                .concatWith(Flux.just(ServerSentEvent.<String>builder()
                        .comment("ended")
                        .build()));
    }

    /**
     * 聊天
     *
     * @param memoryId 内存ID
     * @param message  用户发送的消息
     * @return 聊天结果
     */
    @RequestMapping("/chat")
    public String chat(int memoryId, String message, String token) {
        // 提取请求头的token，用于访问订票系统后端
        log.info("token: {}", token);
        // 放到ThreadLocal中
        TokenContextHolder.setToken(token);
        try {
            return aiFlightHelperService.chat(memoryId, message);
        } finally {
            // 清理ThreadLocal，防止内存泄漏
            TokenContextHolder.clear();
        }
    }

}